/* socket.c : dgnsc socket
 * Copyright (C) 2009-2010 drangon <drangon.zhou@gmail.com>
 * 2009-05
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <dgnsc/socket.h>
#include <dgnsc/common.h>

#ifdef _WIN32
#define FD_SETSIZE	1024
#include <winsock2.h>
#include <windows.h>
#define socklen_t int
#define DGN_SHUT_RDWR	SD_BOTH 
#define IS_ERR_EAGAIN() ( WSAGetLastError() == WSAEWOULDBLOCK )
#else
#include <sys/types.h>
#include <sys/socket.h>
#include <netdb.h>
#include <arpa/inet.h> 
#include <netinet/in.h>
// #include <sys/select.h>
#include <unistd.h>
#include <stdlib.h>
#include <fcntl.h>
#include <string.h>
#include <poll.h>
#include <errno.h>
#define SOCKET int
#define closesocket close
#define INVALID_SOCKET (-1)
#define DGN_SHUT_RDWR	SHUT_RDWR
#define IS_ERR_EAGAIN() ( errno == EAGAIN )
#endif

struct dgn_sock_st
{
	SOCKET sock;
	void * tag;
};

static int set_nonblock( struct dgn_sock_st * sk, int nonblock )
{
#ifdef _WIN32
	int ret;
	unsigned long nb = nonblock ? 1 : 0;
	ret = ioctlsocket( sk->sock, FIONBIO, &nb );
	return ret == SOCKET_ERROR ? -1 : 0;
#else
	int ret, flags;
	flags = fcntl( sk->sock, F_GETFL );
	if( flags < 0 ) {
		return -1;
	}
	if( nonblock ) 
		flags |= O_NONBLOCK;
	else
		flags &= ~O_NONBLOCK;
	ret = fcntl( sk->sock, F_SETFL, flags );
	return ret < 0 ? -1 : 0;
#endif
}

static int sockaddr_aton( const char * host, int port, struct sockaddr_in * addr )
{
#ifdef _WIN32
	// mingw32 does not have getaddrinfo(), but ms2005 have
	struct hostent * hptr = gethostbyname( host );
	if( hptr == NULL || hptr->h_addr_list == NULL || hptr->h_addr_list[0] == NULL )
		return -1;
	memcpy( &(addr->sin_addr), hptr->h_addr_list[0], sizeof(addr->sin_addr) );

#else
	struct addrinfo hint;
	struct addrinfo * res = NULL;
	memset( &hint, 0, sizeof(hint) );
	hint.ai_family = AF_INET;
	hint.ai_socktype = SOCK_STREAM; // enough, useable for DGRAM too
	int ret = getaddrinfo( host, NULL, &hint, &res );
	if( ret != 0 || res == NULL )
		return -1;
	memcpy( addr, res->ai_addr, sizeof( struct sockaddr_in ) );
	if( res != NULL )
		freeaddrinfo( res ), res = NULL;
#endif
	addr->sin_family = AF_INET;
	addr->sin_port = htons( port );
	return 0;
}

static int sockaddr_ntoa( const struct sockaddr_in * addr, char ip[32], int * port )
{
#ifdef _WIN32
	char * tmp = inet_ntoa( addr->sin_addr );
	if( tmp == NULL )
		return -1;
	strcpy( ip, tmp ); // NOTE : ip buf should big enough
#else
	inet_ntop( AF_INET, (void *)&(addr->sin_addr), ip, INET_ADDRSTRLEN );
#endif
	if( port != NULL )
		*port = ntohs( addr->sin_port );
	return 0;
}

int dgn_sock_host_to_ip( const char * host, char ip[32] )
{
	struct sockaddr_in addr;
	if( sockaddr_aton( host, 0, &addr ) < 0 )
		return -1;
	if( sockaddr_ntoa( &addr, ip, NULL ) < 0 )
		return -1;
	return 0;
}

dgn_sock_t dgn_sock_tcp_connect( const char * remote_host, int port, int timeout_ms )
{
	struct sockaddr_in addr;
	if( sockaddr_aton( remote_host, port, &addr ) < 0 )
		return NULL;

	struct dgn_sock_st * sk = (struct dgn_sock_st *)malloc( sizeof(struct dgn_sock_st) );
	sk->tag = NULL;
	sk->sock = socket( AF_INET, SOCK_STREAM, 0 );
	if( sk->sock == INVALID_SOCKET ) {
		dgn_sock_close( sk );
		return NULL;
	}
	if( set_nonblock( sk, 1 ) < 0 ) {
		dgn_sock_close( sk );
		return NULL;
	}

	socklen_t addrlen = sizeof(addr);
	int ret = connect( sk->sock, (struct sockaddr *)&addr, addrlen );
	if( ret >= 0 )
		return sk;

#ifdef _WIN32
	if( ret < 0 && WSAGetLastError() != WSAEWOULDBLOCK )
#else
	if( ret < 0 && errno != EINPROGRESS )
#endif
	{
		dgn_sock_close( sk );
		return NULL;
	}

	if( timeout_ms <= 0 )
		return sk; // sock still in progress

	int ret_evt;
	ret = dgn_sock_poll( sk, DGN_POLLOUT, &ret_evt, timeout_ms );
	if( ret <= 0 ) {
		dgn_sock_close( sk );
		return NULL;
	}

	if( ! dgn_sock_is_connected( sk ) ) {
		dgn_sock_close( sk );
		return NULL;
	}

	return sk; // sock is connected
}

int dgn_sock_is_connected( dgn_sock_t sock )
{
	struct dgn_sock_st * sk = (struct dgn_sock_st *)sock;
#ifdef _WIN32
	fd_set rset, wset, eset;
	struct timeval tv;

	FD_ZERO( &rset );
	FD_ZERO( &wset );
	FD_ZERO( &eset );
	FD_SET( sk->sock, &wset );
	FD_SET( sk->sock, &eset );
	tv.tv_sec = 0;
	tv.tv_usec = 0;

	int ret = select( 0, &rset, &wset, &eset, &tv );
	if( ret > 0 && FD_ISSET( sk->sock, &wset ) )
		return 1; // connected
	if( ret > 0 && FD_ISSET( sk->sock, &eset ) )
		return 0; // connect failed
	return 0; // still in progress

#else
	struct pollfd pfd;
	pfd.fd = sk->sock;
	pfd.revents = 0;
	pfd.events |= POLLOUT;

	int ret = poll( &pfd, 1, 0 );
	if( pfd.revents & POLLOUT ) {
		int err = 0;
		socklen_t len = sizeof(err);
		ret = getsockopt( sk->sock, SOL_SOCKET, SO_ERROR, (void *)&err, &len );
		if( ret == 0 && err == 0 )
			return 1; // connected
		return 0; // failed
	}

	return 0; // still in progress
#endif
}

dgn_sock_t dgn_sock_tcp_svr( const char * listen_host, int port )
{
	struct sockaddr_in addr;
	if( sockaddr_aton( listen_host, port, &addr ) < 0 )
		return NULL;

	struct dgn_sock_st * sk = (struct dgn_sock_st *)malloc( sizeof(struct dgn_sock_st) );
	sk->tag = NULL;
	sk->sock = socket( AF_INET, SOCK_STREAM, 0 );
	if( sk->sock == INVALID_SOCKET ) {
		dgn_sock_close( sk );
		return NULL;
	}
	if( set_nonblock( sk, 1 ) < 0 ) {
		dgn_sock_close( sk );
		return NULL;
	}

#ifndef _WIN32
	// WIN32 need not set reuseaddr, has other means
	int option = 1;
	if( setsockopt( sk->sock, SOL_SOCKET, SO_REUSEADDR, (char *)&option, sizeof(option) ) < 0 ) {
		dgn_sock_close( sk );
		return NULL;
	}
#endif

	socklen_t addrlen = sizeof(addr);
	if( bind( sk->sock, (struct sockaddr *)&addr, addrlen ) < 0 ) {
		dgn_sock_close( sk );
		return NULL;
	}
	if( listen( sk->sock, 7 ) < 0 ) {
		dgn_sock_close( sk );
		return NULL;
	}

	return sk;
}

dgn_sock_t dgn_sock_udp( const char * bind_host, int port )
{
	struct sockaddr_in addr;
	if( sockaddr_aton( bind_host, port, &addr ) < 0 )
		return NULL;

	struct dgn_sock_st * sk = (struct dgn_sock_st *)malloc( sizeof(struct dgn_sock_st) );
	sk->tag = NULL;
	sk->sock = socket( AF_INET, SOCK_DGRAM, 0 );
	if( sk->sock == INVALID_SOCKET ) {
		dgn_sock_close( sk );
		return NULL;
	}

	if( set_nonblock( sk, 1 ) < 0 ) {
		dgn_sock_close( sk );
		return NULL;
	}

#ifndef _WIN32
	// WIN32 need not set reuseaddr, has other means
	int option = 1;
	if( setsockopt( sk->sock, SOL_SOCKET, SO_REUSEADDR, (char *)&option, sizeof(option) ) < 0 ) {
		dgn_sock_close( sk );
		return NULL;
	}
#endif

	socklen_t addrlen = sizeof(addr);
	if( bind( sk->sock, (struct sockaddr *)&addr, addrlen ) < 0 ) {
		dgn_sock_close( sk );
		return NULL;
	}

	return sk;
}

int dgn_sock_local_addr( dgn_sock_t sock, char ip[32], int * port )
{
	struct dgn_sock_st * sk = (struct dgn_sock_st *)sock;
	struct sockaddr_in addr;
	socklen_t len = sizeof(addr);
	int ret = getsockname( sk->sock, (struct sockaddr *)&addr, &len );
	if( ret < 0 )
		return -1;
	return sockaddr_ntoa( &addr, ip, port );
}

int dgn_sock_remote_addr( dgn_sock_t sock, char ip[32], int * port )
{
	struct dgn_sock_st * sk = (struct dgn_sock_st *)sock;
	struct sockaddr_in addr;
	socklen_t len = sizeof(addr);
	int ret = getpeername( sk->sock, (struct sockaddr *)&addr, &len );
	if( ret < 0 )
		return -1;
	return sockaddr_ntoa( &addr, ip, port );
}

dgn_sock_t dgn_sock_accept( dgn_sock_t listen_sock, int timeout_ms )
{
	struct dgn_sock_st * lisk = (struct dgn_sock_st *)listen_sock;
	if( timeout_ms > 0 ) {
		int ret_evt;
		int ret = dgn_sock_poll( lisk, DGN_POLLIN, &ret_evt, timeout_ms );
		if( ret <= 0 ) {
			return NULL;
		}
	}

	SOCKET tmpsk = accept( lisk->sock, NULL, 0 );
	if( tmpsk == INVALID_SOCKET )
		return NULL;

	struct dgn_sock_st * sk = (struct dgn_sock_st *)malloc( sizeof(struct dgn_sock_st) );
	sk->tag = NULL;
	sk->sock = tmpsk;
	if( set_nonblock( sk, 1 ) < 0 ) {
		dgn_sock_close( sk );
		return NULL;
	}

	return sk;
}

void * dgn_sock_get_tag( dgn_sock_t sock )
{
	return ((struct dgn_sock_st *)sock)->tag;
}

void dgn_sock_set_tag( dgn_sock_t sock, void * tag )
{
	((struct dgn_sock_st *)sock)->tag = tag;
	return;
}

int dgn_sock_poll( dgn_sock_t sock, int want_evt, int * ret_evt, int timeout_ms )
{
	struct dgn_sock_st * sk = (struct dgn_sock_st *)sock;
	*ret_evt = 0;
#ifdef _WIN32
	fd_set rset, wset, eset;
	struct timeval tv;

	FD_ZERO( &rset );
	FD_ZERO( &wset );
	FD_ZERO( &eset );
	if( want_evt & DGN_POLLIN )
		FD_SET( sk->sock, &rset );
	if( want_evt & DGN_POLLOUT ) {
		FD_SET( sk->sock, &wset );
		FD_SET( sk->sock, &eset );
	}
	tv.tv_sec = timeout_ms / 1000;
	tv.tv_usec = (timeout_ms % 1000) * 1000;

	int ret = select( 0, &rset, &wset, &eset, &tv );
	if( ret <= 0 )
		return ret;
	if( FD_ISSET( sk->sock, &rset ) )
		*ret_evt |= DGN_POLLIN;
	if( FD_ISSET( sk->sock, &wset ) || FD_ISSET( sk->sock, &eset ) )
		*ret_evt |= DGN_POLLOUT;
	return *ret_evt;

#else
	struct pollfd pfd;
	pfd.fd = sk->sock;
	pfd.events = pfd.revents = 0;
	if( want_evt & DGN_POLLIN )
		pfd.events |= POLLIN;
	if( want_evt & DGN_POLLOUT )
		pfd.events |= POLLOUT;

	int ret = poll( &pfd, 1, timeout_ms );
	if( ret <= 0 )
		return ret;
	if( pfd.revents & POLLIN )
		*ret_evt |= DGN_POLLIN;
	if( pfd.revents & POLLOUT )
		*ret_evt |= DGN_POLLOUT;

	return *ret_evt;
#endif
}

int dgn_sock_pollex( int num, dgn_sock_t * sock, const int * want_evt, int * ret_evt, int timeout_ms )
{
#ifdef _WIN32
	fd_set rset, wset, eset;
	struct timeval tv;
	int i;

	FD_ZERO( &rset );
	FD_ZERO( &wset );
	FD_ZERO( &eset );
	for( i = 0; i < num; ++i ) {
		struct dgn_sock_st * sk = (struct dgn_sock_st *)sock[i];
		if( want_evt[i] & DGN_POLLIN )
			FD_SET( sk->sock, &rset );
		if( want_evt[i] & DGN_POLLOUT ) {
			FD_SET( sk->sock, &wset );
			FD_SET( sk->sock, &eset );
		}
	}
	tv.tv_sec = timeout_ms / 1000;
	tv.tv_usec = (timeout_ms % 1000) * 1000;

	int ret = select( 0, &rset, &wset, &eset, &tv );

	for( i = 0; i < num; ++i ) {
		ret_evt[i] = 0;
	}

	if( ret <= 0 )
		return ret;

	for( i = 0; i < num; ++i ) {
		struct dgn_sock_st * sk = (struct dgn_sock_st *)sock[i];
		if( FD_ISSET( sk->sock, &rset ) )
			ret_evt[i] |= DGN_POLLIN;
		if( FD_ISSET( sk->sock, &wset ) || FD_ISSET( sk->sock, &eset ) )
			ret_evt[i] |= DGN_POLLOUT;
	}
	return ret;

#else
	struct pollfd pfd[1024];
	int i;
	for( i = 0; i < num; ++i ) {
		struct dgn_sock_st * sk = (struct dgn_sock_st *)sock[i];
		pfd[i].fd = sk->sock;
		pfd[i].events = pfd[i].revents = 0;
		if( want_evt[i] & DGN_POLLIN )
			pfd[i].events |= POLLIN;
		if( want_evt[i] & DGN_POLLOUT )
			pfd[i].events |= POLLOUT;
	}

	int ret = poll( pfd, num, timeout_ms );

	for( i = 0; i < num; ++i ) {
		ret_evt[i] = 0;
	}

	if( ret <= 0 )
		return ret;

	for( i = 0; i < num; ++i ) {
		if( pfd[i].revents & POLLIN )
			ret_evt[i] |= DGN_POLLIN;
		if( pfd[i].revents & POLLOUT )
			ret_evt[i] |= DGN_POLLOUT;
	}

	return ret;
#endif
}

int dgn_sock_send( dgn_sock_t sock, const char * buf, int len, int timeout_ms )
{
	int currlen = 0;
	struct dgn_sock_st * sk = (struct dgn_sock_st *)sock;
	int ret = send( sk->sock, buf, len, 0 );
	if( ret == len )
		return ret;
	if( ret < 0 && ! IS_ERR_EAGAIN() )
		return ret;
	if( timeout_ms <= 0 ) {
		return ret < 0 ? 0 : ret;
	}

	if( ret > 0 )
		currlen = ret;

	int64_t now, tm_end;
	now = dgn_time_now();
	tm_end = now + timeout_ms * 1000;

	while( currlen < len ) {
		now = dgn_time_now();
		int64_t tm_diff = tm_end - now;
		if( tm_diff < 0 )
			return currlen;
		if( tm_diff > timeout_ms * 1000 ) {
			// system time has been changed, adjust
			tm_diff = timeout_ms * 1000;
			tm_end = now + tm_diff;
		}

		int ret_evt = 0;
		ret = dgn_sock_poll( sk, DGN_POLLOUT, &ret_evt, tm_diff / 1000 );
		if( ret <= 0 ) {
			return ret < 0 && currlen == 0 ? -1 : currlen;
		}

		ret = send( sk->sock, buf + currlen, len - currlen, 0 );
		if( ret < 0 && ! IS_ERR_EAGAIN() ) {
			return currlen == 0 ? -1 : currlen;
		}
		if( ret > 0 )
			currlen += ret;
	}
	return currlen;
}

int dgn_sock_recv( dgn_sock_t sock, char * buf, int minlen, int maxlen, int timeout_ms )
{
	int currlen = 0;
	struct dgn_sock_st * sk = (struct dgn_sock_st *)sock;
	int ret = recv( sk->sock, buf, maxlen, 0 );
	if( ret >= minlen )
		return ret;
	if( ret == 0 )
		return -1; // peer reset, peer closed, treat as error
	if( ret < 0 && ! IS_ERR_EAGAIN() )
		return ret;
	if( timeout_ms <= 0 ) {
		return ret < 0 ? 0 : ret;
	}

	if( ret > 0 )
		currlen = ret;

	int64_t now, tm_end;
	now = dgn_time_now();
	tm_end = now + timeout_ms * 1000;

	while( currlen < minlen ) {
		now = dgn_time_now();
		int64_t tm_diff = tm_end - now;
		if( tm_diff < 0 )
			return currlen;
		if( tm_diff > timeout_ms * 1000 ) {
			// system time has been changed, adjust
			tm_diff = timeout_ms * 1000;
			tm_end = now + tm_diff;
		}

		int ret_evt = 0;
		ret = dgn_sock_poll( sk, DGN_POLLIN, &ret_evt, tm_diff / 1000 );
		if( ret <= 0 ) {
			return ret < 0 && currlen == 0 ? -1 : currlen;
		}

		ret = recv( sk->sock, buf + currlen, maxlen - currlen, 0 );
		if( ret < 0 && ! IS_ERR_EAGAIN() ) {
			return currlen == 0 ? -1 : currlen;
		}
		if( ret == 0 ) {
			return currlen == 0 ? -1 : currlen;
		}
		if( ret > 0 )
			currlen += ret;
	}
	
	return currlen;
}

int dgn_sock_sendto( dgn_sock_t sock, const char * buf, int len, const char * remote_host, int port )
{
	struct dgn_sock_st * sk = (struct dgn_sock_st *)sock;
	struct sockaddr_in addr;
	if( sockaddr_aton( remote_host, port, &addr ) < 0 )
		return -2;
	socklen_t addrlen = sizeof(addr);
	return sendto( sk->sock, buf, len, 0, (struct sockaddr *)&addr, addrlen );
}

int dgn_sock_recvfrom( dgn_sock_t sock, char * buf, int len, char remote_ip[32], int * port, int timeout_ms )
{
	struct dgn_sock_st * sk = (struct dgn_sock_st *)sock;
	if( timeout_ms > 0 ) {
		int ret_evt;
		int ret = dgn_sock_poll( sk, DGN_POLLIN, &ret_evt, timeout_ms );
		if( ret <= 0 ) {
			return -2;
		}
	}

	struct sockaddr_in addr;
	socklen_t addrlen = sizeof(addr);
	int ret = recvfrom( sk->sock, buf, len, 0, (struct sockaddr *)&addr, &addrlen );

	if( ret > 0 )
		sockaddr_ntoa( &addr, remote_ip, port );

	return ret;
}

void dgn_sock_close( dgn_sock_t sock )
{
	struct dgn_sock_st * sk = (struct dgn_sock_st *)sock;
	if( sk->sock != INVALID_SOCKET ) {
		shutdown( sk->sock, DGN_SHUT_RDWR );
		closesocket( sk->sock );
	}
	free( sk ), sk = NULL;
	return;
}

